import pysrt
import json
import os
import chardet
from transformers import pipeline  # 引入 NLP 模型
import logging

# 加载 GPT 模型
# text_generator = pipeline("text-generation", model="gpt2")
text_generator = pipeline("text-generation", model="gpt2", device=0)  # 使用 GPU

# Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation. 减少输出
logging.getLogger("transformers").setLevel(logging.ERROR)


def parse_subtitles(file_path):
    try:
        with open(file_path, 'rb') as f:
            # 自动检测文件编码
            raw_data = f.read()
            result = chardet.detect(raw_data)
            encoding = result['encoding']

        # 使用检测到的编码读取文件
        subs = pysrt.open(file_path, encoding=encoding)
        subtitles = []
        file_name = os.path.basename(file_path).replace(".srt", "")  # 去掉文件后缀

        for sub in subs:
            text = sub.text.replace("\n", " ").strip()
            try:
                # 使用 GPT 模型生成销售话术
                output_text = text_generator(
                    f"Convert this casual conversation into a sales pitch: {text}",
                    max_new_tokens=50,  # 仅使用 max_new_tokens 控制生成的新文本长度
                    num_return_sequences=1,
                    truncation=True  # 显式启用截断
                )[0]['generated_text']
            except Exception as e:
                print(f"模型生成失败: {e}")
                output_text = "生成失败"  # 生成失败时的默认输出
            # print(output_text)
            try:
                # 切割第一个冒号   Convert this casual conversation into a sales pitch
                output_text_1 = output_text.split(":", 1)[0].strip()
                if output_text_1 == "Convert this casual conversation into a sales pitch":
                    output_text = output_text.split(":", 1)[1].strip()
            except:
                pass

            subtitle_entry = {
                "instruction": "Convert this casual conversation into sales pitch",
                "input": f"{file_name}\n{text}",  # input 格式为 "电影名\n字幕"
                "output": output_text  # output 是生成的销售话术
            }
            subtitles.append(subtitle_entry)
        print(f"{file_path}生成成功")

        return subtitles

    except Exception as e:
        print(f"文件读取或处理失败: {e}")
        return []  # 返回空列表表示处理失败



# 将字幕数据保存为 JSON 文件
def save_to_json(subtitles, output_path):
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(subtitles, f, ensure_ascii=False, indent=4)


# 遍历目录中的所有 SRT 文件并提取内容
def extract_all_subtitles(input_dir, output_json_path):
    all_subtitles = []

    # 遍历目录中的所有文件
    for file_name in os.listdir(input_dir):
        if file_name.endswith('.srt'):  # 只处理 SRT 文件
            file_path = os.path.join(input_dir, file_name)
            subtitles = parse_subtitles(file_path)
            all_subtitles.extend(subtitles)  # 将当前文件的字幕内容添加到总列表中

    # 保存所有字幕到一个 JSON 文件
    save_to_json(all_subtitles, output_json_path)
    print(f"字幕已保存到 {output_json_path}")

#60 80 100 ... 300
for i in range(200, 301, 20):
    # 示例：处理指定目录的字幕文件
    input_subtitle_dir = fr'D:\subtitle\moves\Tomatoes\Tomatoes{i}'
    output_json_path = fr'D:\subtitle\moves\movies_json\Tomatoes{i}.json'
    extract_all_subtitles(input_subtitle_dir, output_json_path)


# # 示例：处理 2022 年每个月的字幕文件
# y = 2022
# for i in range(1, 13):
#     input_subtitle_dir = fr'D:\subtitle\moves\{y}\{i}月'  # 字幕文件夹路径
#
#     if i < 10:
#         output_json_path = fr'D:\subtitle\moves\json\{y}0{i}.json'  # 输出 JSON 文件路径
#     else:
#         output_json_path = fr'D:\subtitle\moves\json\{y}{i}.json'
#
#     extract_all_subtitles(input_subtitle_dir, output_json_path)
